package edu.nf.relaxenjoy.common.recommend;


import ai.djl.Device;
import ai.djl.repository.zoo.Criteria;
import ai.djl.training.util.ProgressBar;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.nio.file.Paths;

/**
 * @author zhu
 * @date 2024-12-20
 **/
public final class SentenceEncoder {

  private static final Logger logger = LoggerFactory.getLogger(SentenceEncoder.class);

  public SentenceEncoder() {}

  public Criteria<String, float[]> criteria() {

    Criteria<String, float[]> criteria =
        Criteria.builder()
            .setTypes(String.class, float[].class)
            .optModelPath(Paths.get("/root/jar/distiluse-base-multilingual-cased-v1.zip"))
            .optTranslator(new SentenceTransTranslator())
            .optEngine("PyTorch") // Use PyTorch engine
            .optDevice(Device.cpu())
            .optProgress(new ProgressBar())
            .build();

    return criteria;
  }
}
